package edu.uw.nemo.esu;

import edu.uw.nemo.io.IntegerArraySerializer;
import edu.uw.nemo.model.Metrics;
import org.apache.log4j.Logger;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.function.FlatMapFunction;
import org.apache.spark.broadcast.Broadcast;
import org.apache.spark.storage.StorageLevel;

import java.io.Serializable;
import java.util.*;

public class ESUEnumerator implements Serializable {

    private static Logger logger = Logger.getLogger(ESUEnumerator.class);

    private final ESUGen generator;

    public ESUEnumerator(edu.uw.nemo.esu.ESUGen ESUGen) {
        this.generator = ESUGen;
    }

    public Metrics enumerateSubgraphs(final Broadcast<Map<Integer, Set<Integer>>> broadcastAdj, JavaRDD<Integer> vertices, final int length, String outputFileName) {
        long start = System.currentTimeMillis();
        logger.debug("******************************************************** starting enumerate subgraph.");
        JavaRDD<int[]> javaRDD = vertices.mapPartitions(new FlatMapFunction<Iterator<Integer>, int[]>() {

            public Iterable<int[]> call(Iterator<Integer> integerIterator) throws Exception {
                Map<Integer, Set<Integer>> adj = broadcastAdj.getValue();
                EnumerationCollector result = new EnumerationCollector();
                while (integerIterator.hasNext()) {
                    Integer node = integerIterator.next();
//                    logger.debug("******************************************************** starting motif exploration from vertex [" + node.intValue() + "].");
                    List<Integer> extension = generator.filterGreater(generator.getNeighbours(adj, node), node);
                    generator.extendSubgraph(adj, makeSubGraph(node), extension, length, result);
//                    logger.debug("******************************************************** finished motif exploration from vertex [" + node.intValue() + "].");
                }
                return result.getMotifs();
            }
        });

//        javaRDD.persist(StorageLevel.MEMORY_AND_DISK());
        javaRDD.persist(StorageLevel.DISK_ONLY());

        logger.debug("******************************************************** enumerate subgraph took " + (System.currentTimeMillis() - start) + " milliseconds.");
        if (outputFileName != null && outputFileName.length() > 0) {
            javaRDD.coalesce(1)
                    .map(new IntegerArraySerializer())
                    .saveAsTextFile(outputFileName);
        }
//        System.out.println("just count = " + javaRDD.count());
        long count = javaRDD.count();
        Metrics metrics = new Metrics();
        metrics.addMetric(Metrics.COUNT, count);

        return metrics;
    }

    public Metrics enumerateSubgraphs2(final Broadcast<Map<Integer, Set<Integer>>> broadcastAdj, JavaRDD<Integer> vertices, final int length, String outputFileName) {
        long start = System.currentTimeMillis();
        logger.debug("******************************************************** starting enumerate subgraph.");
        JavaRDD<int[]> javaRDD = vertices.flatMap(new FlatMapFunction<Integer, int[]>() {

            public Iterable<int[]> call(Integer node) throws Exception {
                Map<Integer, Set<Integer>> adj = broadcastAdj.getValue();
                EnumerationCollector result = new EnumerationCollector();
//                    logger.debug("******************************************************** starting motif exploration from vertex [" + node.intValue() + "].");
                List<Integer> extension = generator.filterGreater(generator.getNeighbours(adj, node), node);
                generator.extendSubgraph(adj, makeSubGraph(node), extension, length, result);
//                    logger.debug("******************************************************** finished motif exploration from vertex [" + node.intValue() + "].");
                return result.getMotifs();
            }
        });

        javaRDD.persist(StorageLevel.MEMORY_AND_DISK());
//        javaRDD.persist(StorageLevel.DISK_ONLY());

        logger.debug("******************************************************** enumerate subgraph took " + (System.currentTimeMillis() - start) + " milliseconds.");
        if (outputFileName != null && outputFileName.length() > 0) {
            javaRDD.coalesce(1)
                    .map(new IntegerArraySerializer())
                    .saveAsTextFile(outputFileName);
        }
//        System.out.println("just count = " + javaRDD.count());
        long count = javaRDD.count();
        Metrics metrics = new Metrics();
        metrics.addMetric(Metrics.COUNT, count);

        return metrics;
    }

    private List<Integer> makeSubGraph(Integer node) {
        List<Integer> subGraph = new ArrayList<Integer>();
        subGraph.add(node);
        return subGraph;
    }

}